from ultralytics import YOLO
from PIL import Image
import numpy as np
from symbol_map import CLASS_NAMES

class YoloNoteDetector:
    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.class_names = CLASS_NAMES

    def detect(self, image_path):
        results = self.model(image_path)
        boxes = results[0].boxes.xyxy.cpu().numpy()
        classes = results[0].boxes.cls.cpu().numpy().astype(int)
        confs = results[0].boxes.conf.cpu().numpy()
        # 返回每个目标的类别名、置信度、位置
        return [
            {
                "class": self.class_names[c],
                "conf": float(confs[i]),
                "bbox": list(map(float, boxes[i]))
            }
            for i, c in enumerate(classes)
        ] 